package main

import (
	"context"
	"database/sql"
	"fmt"
	"gopkg.in/yaml.v3"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

// 数据库配置结构体
type DatabaseConfig struct {
	Username string `yaml:"username"`
	Password string `yaml:"password"`
	Host     string `yaml:"host"`
	Port     int    `yaml:"port"`
	Database string `yaml:"database"`
	Charset  string `yaml:"charset"`
	// 连接池配置
	MaxOpenConns    int           `yaml:"max_open_conns"`
	MaxIdleConns    int           `yaml:"max_idle_conns"`
	ConnMaxLifetime time.Duration `yaml:"conn_max_lifetime"`
	ConnMaxIdleTime time.Duration `yaml:"conn_max_idle_time"`
}

// YAML配置文件结构体
type Config struct {
	MySQL DatabaseConfig `yaml:"mysql"`
}

// 默认配置
func getDefaultConfig() *DatabaseConfig {
	return &DatabaseConfig{
		Username:        "root",
		Password:        "",
		Host:            "localhost",
		Port:            3306,
		Database:        "test",
		Charset:         "utf8mb4",
		MaxOpenConns:    25,
		MaxIdleConns:    10,
		ConnMaxLifetime: 5 * time.Minute,
		ConnMaxIdleTime: time.Minute,
	}
}

// 获取配置文件的完整路径
func getConfigFilePath(filename string) string {
	// 首先尝试当前工作目录
	if _, err := os.Stat(filename); err == nil {
		return filename
	}

	// 然后尝试程序所在目录
	execPath, err := os.Executable()
	if err != nil {
		log.Printf("Failed to get executable path: %v", err)
		return filename
	}

	execDir := filepath.Dir(execPath)
	configPath := filepath.Join(execDir, filename)

	log.Printf("Looking for config file at: %s", configPath)
	return configPath
}

// 从YAML配置文件加载配置
func loadConfigFromYAML(filename string) (*DatabaseConfig, error) {
	configPath := getConfigFilePath(filename)

	// 检查文件是否存在
	if _, err := os.Stat(configPath); os.IsNotExist(err) {
		return nil, fmt.Errorf("config file %s does not exist", configPath)
	}

	// 读取文件内容
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read config file: %w", err)
	}

	// 解析YAML
	var config Config
	if err := yaml.Unmarshal(data, &config); err != nil {
		return nil, fmt.Errorf("failed to parse YAML config: %w", err)
	}

	// 设置默认值（如果配置文件中没有指定）
	dbConfig := &config.MySQL
	if dbConfig.Host == "" {
		dbConfig.Host = "localhost"
	}
	if dbConfig.Port == 0 {
		dbConfig.Port = 3306
	}
	if dbConfig.Username == "" {
		dbConfig.Username = "root"
	}
	if dbConfig.Database == "" {
		dbConfig.Database = "test"
	}
	if dbConfig.Charset == "" {
		dbConfig.Charset = "utf8mb4"
	}
	if dbConfig.MaxOpenConns == 0 {
		dbConfig.MaxOpenConns = 25
	}
	if dbConfig.MaxIdleConns == 0 {
		dbConfig.MaxIdleConns = 10
	}
	if dbConfig.ConnMaxLifetime == 0 {
		dbConfig.ConnMaxLifetime = 5 * time.Minute
	}
	if dbConfig.ConnMaxIdleTime == 0 {
		dbConfig.ConnMaxIdleTime = time.Minute
	}

	return dbConfig, nil
}

// 从环境变量加载配置
func loadConfigFromEnv() *DatabaseConfig {
	config := getDefaultConfig()

	if username := os.Getenv("MYSQL_USERNAME"); username != "" {
		config.Username = username
	}
	if password := os.Getenv("MYSQL_PASSWORD"); password != "" {
		config.Password = password
	}
	if host := os.Getenv("MYSQL_HOST"); host != "" {
		config.Host = host
	}
	if portStr := os.Getenv("MYSQL_PORT"); portStr != "" {
		if port, err := strconv.Atoi(portStr); err == nil {
			config.Port = port
		}
	}
	if database := os.Getenv("MYSQL_DATABASE"); database != "" {
		config.Database = database
	}
	if charset := os.Getenv("MYSQL_CHARSET"); charset != "" {
		config.Charset = charset
	}

	// 连接池配置
	if maxOpenStr := os.Getenv("MYSQL_MAX_OPEN_CONNS"); maxOpenStr != "" {
		if maxOpen, err := strconv.Atoi(maxOpenStr); err == nil && maxOpen > 0 {
			config.MaxOpenConns = maxOpen
		}
	}
	if maxIdleStr := os.Getenv("MYSQL_MAX_IDLE_CONNS"); maxIdleStr != "" {
		if maxIdle, err := strconv.Atoi(maxIdleStr); err == nil && maxIdle > 0 {
			config.MaxIdleConns = maxIdle
		}
	}
	if maxLifetimeStr := os.Getenv("MYSQL_CONN_MAX_LIFETIME"); maxLifetimeStr != "" {
		if maxLifetime, err := time.ParseDuration(maxLifetimeStr); err == nil {
			config.ConnMaxLifetime = maxLifetime
		}
	}
	if maxIdleTimeStr := os.Getenv("MYSQL_CONN_MAX_IDLE_TIME"); maxIdleTimeStr != "" {
		if maxIdleTime, err := time.ParseDuration(maxIdleTimeStr); err == nil {
			config.ConnMaxIdleTime = maxIdleTime
		}
	}

	return config
}

// 加载配置（优先环境变量，然后是YAML文件，最后是默认值）
func loadConfig() *DatabaseConfig {
	// 打印调试信息
	if wd, err := os.Getwd(); err == nil {
		log.Printf("Current working directory: %s", wd)
	}
	if execPath, err := os.Executable(); err == nil {
		execDir := filepath.Dir(execPath)
		log.Printf("Executable directory: %s", execDir)
	}

	// 首先尝试从环境变量加载
	envConfig := loadConfigFromEnv()

	// 检查环境变量是否提供了完整的配置（至少要有用户名）
	if envConfig.Username != "root" || envConfig.Password != "" ||
		os.Getenv("MYSQL_USERNAME") != "" || os.Getenv("MYSQL_PASSWORD") != "" {
		log.Println("Using configuration from environment variables")
		return envConfig
	}

	// 尝试从YAML文件加载
	yamlConfig, err := loadConfigFromYAML("env.yml")
	if err != nil {
		log.Printf("Failed to load YAML config: %v", err)
		log.Println("Using default configuration")
		return getDefaultConfig()
	}

	log.Println("Using configuration from env.yml file")
	return yamlConfig
}

// 全局数据库连接池和配置
var (
	db     *sql.DB
	config *DatabaseConfig
)

// 构建数据库连接字符串
func getDSN(cfg *DatabaseConfig) string {
	return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?parseTime=true&charset=%s",
		cfg.Username, cfg.Password, cfg.Host, cfg.Port, cfg.Database, cfg.Charset)
}

// 初始化数据库连接池
func initDB(cfg *DatabaseConfig) error {
	dsn := getDSN(cfg)
	var err error
	db, err = sql.Open("mysql", dsn)
	if err != nil {
		return fmt.Errorf("mysql数据库连接失败: %w", err)
	}

	// 配置连接池
	db.SetMaxOpenConns(cfg.MaxOpenConns)
	db.SetMaxIdleConns(cfg.MaxIdleConns)
	db.SetConnMaxLifetime(cfg.ConnMaxLifetime)
	db.SetConnMaxIdleTime(cfg.ConnMaxIdleTime)

	// 测试连接
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		return fmt.Errorf("failed to ping database: %w", err)
	}

	log.Printf("Database connection pool initialized successfully")
	log.Printf("Connected to: %s@%s:%d/%s", cfg.Username, cfg.Host, cfg.Port, cfg.Database)
	return nil
}

// 关闭数据库连接
func closeDB() {
	if db != nil {
		db.Close()
	}
}

// 验证表名和列名是否安全（简单的白名单验证）
func isValidIdentifier(name string) bool {
	if name == "" {
		return false
	}
	// 只允许字母、数字、下划线，且不能以数字开头
	for i, r := range name {
		if i == 0 && (r >= '0' && r <= '9') {
			return false
		}
		if !((r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') ||
			(r >= '0' && r <= '9') || r == '_') {
			return false
		}
	}
	return true
}

// 验证多个标识符（用逗号分隔），允许通配符*
func areValidIdentifiers(identifiers string) bool {
	// 特殊处理：单独的通配符*是允许的
	if strings.TrimSpace(identifiers) == "*" {
		return true
	}
	
	// 处理多列情况
	names := strings.Split(identifiers, ",")
	for _, name := range names {
		name = strings.TrimSpace(name)
		if !isValidIdentifier(name) {
			return false
		}
	}
	return true
}

// 创建表的工具处理函数
func createTableHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	columns, err := request.RequireString("columns")
	if err != nil {
		return mcp.NewToolResultError("columns parameter is required"), nil
	}

	// 验证表名安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}

	// 使用参数化查询（虽然表名无法参数化，但我们已经验证了安全性）
	query := fmt.Sprintf("CREATE TABLE IF NOT EXISTS `%s` (%s)", tableName, columns)

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	_, err = db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to create table: %v", err)), nil
	}

	return mcp.NewToolResultText(fmt.Sprintf("Table %s created successfully", tableName)), nil
}

// 插入数据的工具处理函数
func insertDataHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	columns, err := request.RequireString("columns")
	if err != nil {
		return mcp.NewToolResultError("columns parameter is required"), nil
	}

	values, err := request.RequireString("values")
	if err != nil {
		return mcp.NewToolResultError("values parameter is required"), nil
	}

	// 验证标识符安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}
	if !areValidIdentifiers(columns) {
		return mcp.NewToolResultError("invalid column names"), nil
	}

	// 构建查询语句
	query := fmt.Sprintf("INSERT INTO `%s` (%s) VALUES (%s)", tableName, columns, values)

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	result, err := db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to insert data: %v", err)), nil
	}

	rowsAffected, _ := result.RowsAffected()
	return mcp.NewToolResultText(fmt.Sprintf("%d rows inserted successfully", rowsAffected)), nil
}

// 查询数据的工具处理函数
func selectDataHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	columns, err := request.RequireString("columns")
	if err != nil {
		return mcp.NewToolResultError("columns parameter is required"), nil
	}

	// 获取可选参数 condition（参考原始代码的处理方式）
	condition, err := request.RequireString("condition")
	if err != nil {
		// 如果参数不存在或获取失败，设为空字符串
		condition = ""
	}

	// 验证标识符安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}
	if !areValidIdentifiers(columns) {
		return mcp.NewToolResultError("invalid column names"), nil
	}

	var query string
	if condition != "" {
		query = fmt.Sprintf("SELECT %s FROM `%s` WHERE %s", columns, tableName, condition)
	} else {
		query = fmt.Sprintf("SELECT %s FROM `%s`", columns, tableName)
	}

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	rows, err := db.QueryContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to query data: %v", err)), nil
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to get columns: %v", err)), nil
	}

	var results []string
	for rows.Next() {
		values := make([]interface{}, len(cols))
		valuePtrs := make([]interface{}, len(cols))
		for i := range cols {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("failed to scan row: %v", err)), nil
		}

		var rowResult []string
		for i, col := range cols {
			var v string
			if val := values[i]; val != nil {
				if b, ok := val.([]byte); ok {
					v = string(b)
				} else {
					v = fmt.Sprintf("%v", val)
				}
			} else {
				v = "NULL"
			}
			rowResult = append(rowResult, fmt.Sprintf("%s: %s", col, v))
		}
		results = append(results, strings.Join(rowResult, ", "))
	}

	if err := rows.Err(); err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("rows iteration error: %v", err)), nil
	}

	if len(results) == 0 {
		return mcp.NewToolResultText("No data found"), nil
	}

	return mcp.NewToolResultText(strings.Join(results, "\n")), nil
}

// 更新数据的工具处理函数
func updateDataHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	setClause, err := request.RequireString("set_clause")
	if err != nil {
		return mcp.NewToolResultError("set_clause parameter is required"), nil
	}

	condition, err := request.RequireString("condition")
	if err != nil {
		return mcp.NewToolResultError("condition parameter is required"), nil
	}

	// 验证表名安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}

	query := fmt.Sprintf("UPDATE `%s` SET %s WHERE %s", tableName, setClause, condition)

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	result, err := db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to update data: %v", err)), nil
	}

	rowsAffected, _ := result.RowsAffected()
	return mcp.NewToolResultText(fmt.Sprintf("%d rows updated successfully", rowsAffected)), nil
}

// 删除数据的工具处理函数
func deleteDataHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	condition, err := request.RequireString("condition")
	if err != nil {
		return mcp.NewToolResultError("condition parameter is required"), nil
	}

	// 验证表名安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}

	query := fmt.Sprintf("DELETE FROM `%s` WHERE %s", tableName, condition)

	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	result, err := db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to delete data: %v", err)), nil
	}

	rowsAffected, _ := result.RowsAffected()
	return mcp.NewToolResultText(fmt.Sprintf("%d rows deleted successfully", rowsAffected)), nil
}

// 创建索引的工具处理函数
func createIndexHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	indexName, err := request.RequireString("index_name")
	if err != nil {
		return mcp.NewToolResultError("index_name parameter is required"), nil
	}

	columns, err := request.RequireString("columns")
	if err != nil {
		return mcp.NewToolResultError("columns parameter is required"), nil
	}

	// 验证标识符安全性
	if !isValidIdentifier(tableName) || !isValidIdentifier(indexName) {
		return mcp.NewToolResultError("invalid table or index name"), nil
	}
	if !areValidIdentifiers(columns) {
		return mcp.NewToolResultError("invalid column names"), nil
	}

	query := fmt.Sprintf("CREATE INDEX `%s` ON `%s` (%s)", indexName, tableName, columns)

	ctx, cancel := context.WithTimeout(ctx, 60*time.Second) // 索引创建可能需要更长时间
	defer cancel()

	_, err = db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to create index: %v", err)), nil
	}

	return mcp.NewToolResultText(fmt.Sprintf("Index %s created successfully on table %s",
		indexName, tableName)), nil
}

// 分析表的工具处理函数
func analyzeTableHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("table_name parameter is required"), nil
	}

	// 验证表名安全性
	if !isValidIdentifier(tableName) {
		return mcp.NewToolResultError("invalid table name"), nil
	}

	query := fmt.Sprintf("ANALYZE TABLE `%s`", tableName)

	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	_, err = db.ExecContext(ctx, query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("failed to analyze table: %v", err)), nil
	}

	return mcp.NewToolResultText(fmt.Sprintf("Table %s analyzed successfully", tableName)), nil
}

func main() {
	// 加载配置
	config = loadConfig()

	// 打印配置信息（密码部分隐藏）
	log.Printf("Loading MySQL MCP Server configuration:")
	log.Printf("  Host: %s:%d", config.Host, config.Port)
	log.Printf("  Username: %s", config.Username)
	log.Printf("  Database: %s", config.Database)
	log.Printf("  Charset: %s", config.Charset)
	log.Printf("  Max Open Connections: %d", config.MaxOpenConns)
	log.Printf("  Max Idle Connections: %d", config.MaxIdleConns)

	// 初始化数据库连接池
	if err := initDB(config); err != nil {
		log.Fatalf("Failed to initialize database: %v", err)
	}
	defer closeDB()

	// 创建MCP服务器
	s := server.NewMCPServer(
		"MySQL MCP Server",
		"1.0.0",
		server.WithToolCapabilities(true),
		server.WithRecovery(),
	)

	// 注册工具
	createTableTool := mcp.NewTool("create_table",
		mcp.WithDescription("Create a new table in the database"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to create"),
		),
		mcp.WithString("columns",
			mcp.Required(),
			mcp.Description("The column definitions for the table"),
		),
	)
	s.AddTool(createTableTool, createTableHandler)

	insertDataTool := mcp.NewTool("insert_data",
		mcp.WithDescription("Insert data into a table"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to insert data into"),
		),
		mcp.WithString("columns",
			mcp.Required(),
			mcp.Description("The columns to insert data into"),
		),
		mcp.WithString("values",
			mcp.Required(),
			mcp.Description("The values to insert"),
		),
	)
	s.AddTool(insertDataTool, insertDataHandler)

	selectDataTool := mcp.NewTool("select_data",
		mcp.WithDescription("Select data from a table"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to select data from"),
		),
		mcp.WithString("columns",
			mcp.Required(),
			mcp.Description("The columns to select"),
		),
		mcp.WithString("condition",
			mcp.Description("The optional WHERE condition"),
		),
	)
	s.AddTool(selectDataTool, selectDataHandler)

	updateDataTool := mcp.NewTool("update_data",
		mcp.WithDescription("Update data in a table"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to update data in"),
		),
		mcp.WithString("set_clause",
			mcp.Required(),
			mcp.Description("The SET clause for the update"),
		),
		mcp.WithString("condition",
			mcp.Required(),
			mcp.Description("The WHERE condition for the update"),
		),
	)
	s.AddTool(updateDataTool, updateDataHandler)

	deleteDataTool := mcp.NewTool("delete_data",
		mcp.WithDescription("Delete data from a table"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to delete data from"),
		),
		mcp.WithString("condition",
			mcp.Required(),
			mcp.Description("The WHERE condition for the delete"),
		),
	)
	s.AddTool(deleteDataTool, deleteDataHandler)

	createIndexTool := mcp.NewTool("create_index",
		mcp.WithDescription("Create an index on a table"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to create the index on"),
		),
		mcp.WithString("index_name",
			mcp.Required(),
			mcp.Description("The name of the index"),
		),
		mcp.WithString("columns",
			mcp.Required(),
			mcp.Description("The columns to create the index on"),
		),
	)
	s.AddTool(createIndexTool, createIndexHandler)

	analyzeTableTool := mcp.NewTool("analyze_table",
		mcp.WithDescription("Analyze a table for optimization"),
		mcp.WithString("table_name",
			mcp.Required(),
			mcp.Description("The name of the table to analyze"),
		),
	)
	s.AddTool(analyzeTableTool, analyzeTableHandler)

	// 启动服务器
	log.Println("Starting MySQL MCP Server...")
	if err := server.ServeStdio(s); err != nil {
		log.Fatalf("Server error: %v", err)
	}
}
